"""
WebSocket实时数据路由
"""
import asyncio
import json
from datetime import datetime
from typing import Optional
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query

from auth.jwt_handler import verify_token
from models.schemas import WebSocketMessage, WebSocketStatusData, CurrentTask


router = APIRouter(tags=["实时数据"])


class ConnectionManager:
    """WebSocket连接管理器"""
    
    def __init__(self):
        self.active_connections: dict[int, list[WebSocket]] = {}
    
    async def connect(self, websocket: WebSocket, workstation_id: int):
        """连接WebSocket"""
        await websocket.accept()
        if workstation_id not in self.active_connections:
            self.active_connections[workstation_id] = []
        self.active_connections[workstation_id].append(websocket)
        print(f"✅ WebSocket连接建立: 工位ID={workstation_id}")
    
    def disconnect(self, websocket: WebSocket, workstation_id: int):
        """断开WebSocket连接"""
        if workstation_id in self.active_connections:
            self.active_connections[workstation_id].remove(websocket)
            if not self.active_connections[workstation_id]:
                del self.active_connections[workstation_id]
        print(f"❌ WebSocket连接断开: 工位ID={workstation_id}")
    
    async def send_message(self, message: dict, workstation_id: int):
        """向指定工位的所有连接发送消息"""
        if workstation_id in self.active_connections:
            disconnected = []
            for connection in self.active_connections[workstation_id]:
                try:
                    await connection.send_json(message)
                except Exception as e:
                    print(f"❌ 发送消息失败: {e}")
                    disconnected.append(connection)
            
            # 清理断开的连接
            for conn in disconnected:
                self.disconnect(conn, workstation_id)


manager = ConnectionManager()


async def generate_mock_data(workstation_id: int) -> dict:
    """生成模拟实时数据"""
    import random
    
    # 模拟不同的消息类型
    message_types = ["status_update", "alert", "task_complete"]
    message_type = random.choice(message_types)
    
    now = datetime.now().isoformat()
    
    if message_type == "status_update":
        # 状态更新消息
        data = WebSocketStatusData(
            workstation_id=workstation_id,
            status=random.choice(["idle", "running", "error"]),
            cpu_usage=random.randint(30, 90),
            memory_usage=random.randint(40, 85),
            task_progress=random.randint(0, 100),
            running_time=random.randint(0, 36000),
            completed_tasks=random.randint(100, 200),
            error_count=random.randint(0, 10),
            current_task=CurrentTask(
                id=f"task_{random.randint(1, 100):03d}",
                name=random.choice(["装配任务", "检测任务", "搬运任务", "焊接任务"]),
                progress=random.randint(0, 100)
            )
        )
        
        message = WebSocketMessage(
            type="status_update",
            timestamp=now,
            data=data.model_dump()
        )
    
    elif message_type == "alert":
        # 告警消息
        alerts = [
            {"level": "warning", "message": "温度超过阈值", "details": {"current_temp": 85, "threshold": 80}},
            {"level": "warning", "message": "内存使用率过高", "details": {"current": 90, "threshold": 85}},
            {"level": "error", "message": "设备通信异常", "details": {"device": "PLC-001", "error": "timeout"}},
            {"level": "info", "message": "任务队列已满", "details": {"queue_size": 100, "max_size": 100}}
        ]
        alert = random.choice(alerts)
        
        message = WebSocketMessage(
            type="alert",
            timestamp=now,
            data={
                "workstation_id": workstation_id,
                "level": alert["level"],
                "message": alert["message"],
                "details": alert["details"]
            }
        )
    
    else:  # task_complete
        # 任务完成消息
        message = WebSocketMessage(
            type="task_complete",
            timestamp=now,
            data={
                "workstation_id": workstation_id,
                "task_id": f"task_{random.randint(1, 100):03d}",
                "task_name": random.choice(["装配任务", "检测任务", "搬运任务"]),
                "status": "success",
                "duration": random.randint(60, 600),
                "result": {"quality": "合格", "count": random.randint(1, 10)}
            }
        )
    
    return message.model_dump()


@router.websocket("/workstations/{id}/realtime")
async def workstation_realtime(
    websocket: WebSocket,
    id: int,
    token: str = Query(..., description="JWT Token")
):
    """
    通过WebSocket推送工位实时状态数据
    
    连接URL: ws://localhost:8000/api/v1/drivionExe/workstations/{id}/realtime?token={jwt_token}
    
    消息类型：
    - status_update: 工位状态更新
    - alert: 告警信息
    - task_complete: 任务完成通知
    - error: 错误信息
    """
    # 验证token
    payload = verify_token(token)
    if payload is None:
        await websocket.close(code=1008, reason="Token无效或已过期")
        return
    
    # 建立连接
    await manager.connect(websocket, id)
    
    try:
        # 发送欢迎消息
        welcome_message = {
            "type": "connected",
            "timestamp": datetime.now().isoformat(),
            "data": {
                "workstation_id": id,
                "message": f"已连接到工位 {id} 的实时数据流"
            }
        }
        await websocket.send_json(welcome_message)
        
        # 启动数据推送任务
        push_task = asyncio.create_task(push_realtime_data(websocket, id))
        
        # 保持连接，接收客户端消息（如果有）
        while True:
            try:
                data = await websocket.receive_text()
                # 可以处理客户端发送的消息（如控制指令等）
                print(f"📨 收到客户端消息: {data}")
            except WebSocketDisconnect:
                break
    
    except Exception as e:
        print(f"❌ WebSocket错误: {e}")
    
    finally:
        # 取消推送任务
        if 'push_task' in locals():
            push_task.cancel()
        # 断开连接
        manager.disconnect(websocket, id)


async def push_realtime_data(websocket: WebSocket, workstation_id: int):
    """推送实时数据"""
    try:
        while True:
            # 每2-5秒推送一次数据
            await asyncio.sleep(3)
            
            # 生成并发送模拟数据
            message = await generate_mock_data(workstation_id)
            await websocket.send_json(message)
    
    except asyncio.CancelledError:
        print(f"⏹️ 停止推送数据: 工位ID={workstation_id}")
    except Exception as e:
        print(f"❌ 推送数据错误: {e}")
